#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
联邦学习与工业低碳优化系统集成示例
"""

from modules.federated.server import FederatedServer
from modules.federated.client import FederatedClient
from modules.federated.model import IndustrialFederatedModel
from modules.federated.config import FederatedConfig
from modules.federated.aggregator import ByzantineRobustAggregator
from modules.industrial import IndustrialLowCarbon
import torch
import numpy as np

def main():
    # 1. 初始化配置
    config = FederatedConfig("example_config.yaml")
    
    # 2. 创建联邦模型 (与现有工业模型兼容)
    industrial_model = IndustrialLowCarbon()
    fed_model = IndustrialFederatedModel(
        input_size=industrial_model.input_dim,
        output_size=industrial_model.output_dim
    )
    
    # 3. 初始化服务器
    server = FederatedServer(fed_model, config)
    aggregator = ByzantineRobustAggregator(config)
    
    # 4. 模拟客户端注册
    for i in range(5):
        config.register_client(f"factory_{i}", {
            "location": f"region_{i%3}",
            "type": "steel" if i < 3 else "chemical"
        })
    
    # 5. 联邦训练循环
    for round in range(config.config['federated']['rounds']):
        print(f"\n=== Round {round + 1} ===")
        
        # 选择客户端
        clients = config.get_client_selection(round)
        print(f"Selected clients: {clients}")
        
        # 模拟客户端训练
        client_updates = []
        for client_id in clients:
            client_model = IndustrialFederatedModel.from_parameters(
                fed_model.get_parameters()
            )
            client = FederatedClient(client_id, client_model, config)
            
            # 模拟训练数据 (实际应用中替换为真实数据)
            train_data = torch.randn(100, industrial_model.input_dim)
            train_labels = torch.randint(0, industrial_model.output_dim, (100,))
            
            # 本地训练并收集更新
            update = client.train([(train_data, train_labels)], epochs=2)
            client_updates.append(update)
        
        # 安全聚合
        global_update = aggregator.aggregate(client_updates)
        server.aggregate([global_update])
        
        # 评估
        test_data = torch.randn(50, industrial_model.input_dim)
        test_labels = torch.randint(0, industrial_model.output_dim, (50,))
        metrics = server.evaluate_global_model([(test_data, test_labels)])
        print(f"Global model accuracy: {metrics['accuracy']:.2%}")

if __name__ == "__main__":
    main()